package org.os.framework.crypto.engine;

import java.math.BigInteger;
import java.security.SecureRandom;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.params.ECKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;
import org.bouncycastle.crypto.signers.ECDSASigner;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.encoders.Hex;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;

/**
 * @author
 */
public class Sm2EcdsaSigner extends ECDSASigner {
    private final Digest digest;
    private ECKeyParameters key;
    private SecureRandom random;

    public Sm2EcdsaSigner(Digest digest) {
        this.digest = digest;
    }

    private byte[] hashMessageIfNeeded(byte[] message) {
        if (this.digest != null) {
            this.digest.reset();
            this.digest.update(message, 0, message.length);
            byte[] out = new byte[digest.getDigestSize()];
            int len = this.digest.doFinal(out, 0);
            byte[] o = new byte[len];
            System.arraycopy(out, 0, o, 0, len);
            return o;
        }
        return message;
    }

    @Override
    public void init(boolean forSigning, CipherParameters param) {
        if (forSigning) {
            if (param instanceof ParametersWithRandom) {
                ParametersWithRandom rParam = (ParametersWithRandom) param;
                this.key = (ECPrivateKeyParameters) rParam.getParameters();
                this.random = rParam.getRandom();
            } else {
                this.key = (ECPrivateKeyParameters) param;
                this.random = new SecureRandom();
            }
        } else {
            this.key = (ECPublicKeyParameters) param;
        }
    }

    @Override
    public BigInteger[] generateSignature(byte[] message) {
        ECPrivateKeyParameters priKey = (ECPrivateKeyParameters) this.key;
        message = hashMessageIfNeeded(message);
        BigInteger m = new BigInteger(Hex.toHexString(message), 16);
        BigInteger n = priKey.getParameters().getN();
        BigInteger d = priKey.getD();
        RandomCalculator kCalculator = new RandomCalculator(n, this.random);
        BigInteger r, s, k;
        while (true) {
            k = kCalculator.nextK();
            ECPoint kg = (priKey.getParameters().getG().multiply(k)).normalize();
            r = (m.add(kg.getXCoord().toBigInteger())).mod(n);
            if (r.compareTo(BigInteger.ZERO) == 0) {
                continue;
            }
            if ((r.add(k)).compareTo(n) == 0) {
                continue;
            }
            BigInteger t1 = (BigInteger.ONE.add(d)).modInverse(n);
            BigInteger t2 = (k.subtract(r.multiply(d))).mod(n);
            s = (t1.multiply(t2)).mod(n);
            if (Hex.toHexString(r.toByteArray()).length() != 64 || Hex.toHexString(s.toByteArray()).length() != 64) {
                continue;
            }
            if (s.compareTo(BigInteger.ZERO) != 0) {
                break;
            }
        }
        return new BigInteger[] { r, s };
    }

    @Override
    public boolean verifySignature(byte[] message, BigInteger r, BigInteger s) {
        ECPublicKeyParameters pubKey = (ECPublicKeyParameters) this.key;
        message = hashMessageIfNeeded(message);
        BigInteger m = new BigInteger(Hex.toHexString(message), 16);
        BigInteger n = pubKey.getParameters().getN();
        if (r.compareTo(ONE) < 0 || r.compareTo(n) >= 0) {
            return false;
        }
        // s in the range [1,n-1]
        if (s.compareTo(ONE) < 0 || s.compareTo(n) >= 0) {
            return false;
        }
        final BigInteger t = (r.add(s)).mod(n);
        if (t.compareTo(BigInteger.ZERO) == 0) {
            return false;
        }
        ECPoint g = pubKey.getParameters().getG().normalize();
        ECPoint q = pubKey.getQ().normalize();
        ECPoint gs = g.multiply(s).normalize();
        ECPoint qt = q.multiply(t).normalize();
        ECPoint q1 = gs.add(qt).normalize();
        BigInteger r1 = (m.add(q1.getXCoord().toBigInteger())).mod(n);
        return r1.compareTo(r) == 0;
    }

}
